"""
  This script provides an example to wrap UER-py for NER.
"""
import random
import argparse
import json
import torch
import torch.nn as nn
from uer.layers.embeddings import *
from uer.encoders.bert_encoder import *
from uer.encoders.rnn_encoder import *
from uer.encoders.birnn_encoder import *
from uer.encoders.cnn_encoder import *
from uer.encoders.attn_encoder import *
from uer.encoders.gpt_encoder import *
from uer.encoders.mixed_encoder import *
from uer.utils.config import load_hyperparam
from uer.utils.optimizers import *
from uer.utils.constants import *
from uer.utils.vocab import Vocab
from uer.utils.seed import set_seed
from uer.utils.tokenizer import *
from uer.model_saver import save_model
from run_classifier import build_optimizer, load_or_initialize_parameters


class NerTagger(nn.Module):
    def __init__(self, args):
        super(NerTagger, self).__init__()
        self.embedding = globals()[args.embedding.capitalize() + "Embedding"](args, len(args.tokenizer.vocab))
        self.encoder = globals()[args.encoder.capitalize() + "Encoder"](args)
        self.labels_num = args.labels_num
        self.output_layer = nn.Linear(args.hidden_size, self.labels_num)

    def forward(self, src, tgt, seg):
        """
        Args:
            src: [batch_size x seq_length]
            tgt: [batch_size x seq_length]
            seg: [batch_size x seq_length]
        Returns:
            loss: Sequence labeling loss.
            logits: Output logits.
        """
        # Embedding.
        emb = self.embedding(src, seg)
        # Encoder.
        output = self.encoder(emb, seg)
        # Target.
        logits = self.output_layer(output).contiguous().view(-1, self.labels_num)

        if tgt is not None:
            tgt = tgt.contiguous().view(-1,1)
            one_hot = torch.zeros(tgt.size(0), self.labels_num). \
                      to(torch.device(tgt.device)). \
                      scatter_(1, tgt, 1.0)

            numerator = -torch.sum(nn.LogSoftmax(dim=-1)(logits) * one_hot, 1)
            
            tgt = tgt.contiguous().view(-1)
            tgt_mask = (tgt<self.labels_num-1).float().to(torch.device(tgt.device))

            numerator = torch.sum(tgt_mask * numerator)
            denominator = torch.sum(tgt_mask) + 1e-6
            loss = numerator / denominator
            return loss, logits
        else:
            return None, logits


def read_dataset(args, path):
    dataset, columns = [], {}
    with open(path, mode="r", encoding="utf-8") as f:
        for line_id, line in enumerate(f):
            if line_id == 0:
                for i, column_name in enumerate(line.strip().split("\t")):
                    columns[column_name] = i
                continue
            line = line.strip().split('\t')
            labels = line[columns["label"]]
            tgt = [args.l2i[l] for l in labels.split(" ")]

            text_a = line[columns["text_a"]]
            src = args.tokenizer.convert_tokens_to_ids(args.tokenizer.tokenize(text_a))
            seg = [1] * len(src)

            if len(src) > args.seq_length:
                src = src[:args.seq_length]
                tgt = tgt[:args.seq_length]
                seg = seg[:args.seq_length]
            while len(src) < args.seq_length:
                src.append(0)
                tgt.append(args.labels_num-1)
                seg.append(0)
            dataset.append([src, tgt, seg])

    return dataset


def batch_loader(batch_size, src, tgt, seg):
    instances_num = src.size()[0]
    for i in range(instances_num // batch_size):
        src_batch = src[i*batch_size: (i+1)*batch_size, :]
        tgt_batch = tgt[i*batch_size: (i+1)*batch_size, :]
        seg_batch = seg[i*batch_size: (i+1)*batch_size, :]
        yield src_batch, tgt_batch, seg_batch
    if instances_num > instances_num // batch_size * batch_size:
        src_batch = src[instances_num//batch_size*batch_size:, :]
        tgt_batch = tgt[instances_num//batch_size*batch_size:, :]
        seg_batch = seg[instances_num//batch_size*batch_size:, :]
        yield src_batch, tgt_batch, seg_batch


def train(args, model, optimizer, scheduler, src_batch, tgt_batch, seg_batch):
    model.zero_grad()

    src_batch = src_batch.to(args.device)
    tgt_batch = tgt_batch.to(args.device)
    seg_batch = seg_batch.to(args.device)

    loss, _ = model(src_batch, tgt_batch, seg_batch)
    if torch.cuda.device_count() > 1:
        loss = torch.mean(loss)

    if args.fp16:
        with amp.scale_loss(loss, optimizer) as scaled_loss:
            scaled_loss.backward()
    else:
        loss.backward()

    optimizer.step()
    scheduler.step()

    return loss


def evaluate(args, dataset):
    src = torch.LongTensor([sample[0] for sample in dataset])
    tgt = torch.LongTensor([sample[1] for sample in dataset])
    seg = torch.LongTensor([sample[2] for sample in dataset])

    instances_num = src.size(0)
    batch_size = args.batch_size

    correct, gold_entities_num, pred_entities_num = 0, 0, 0

    args.model.eval()

    for i, (src_batch, tgt_batch, seg_batch) in enumerate(batch_loader(batch_size, src, tgt, seg)):
        src_batch = src_batch.to(args.device)
        tgt_batch = tgt_batch.to(args.device)
        seg_batch = seg_batch.to(args.device)
        loss, logits = args.model(src_batch, tgt_batch, seg_batch)

        pred = logits.argmax(dim=-1)
        gold = tgt_batch.contiguous().view(-1,1)

        for j in range(gold.size()[0]):
            if gold[j].item() in args.begin_ids:
                gold_entities_num += 1

        for j in range(pred.size()[0]):
            if pred[j].item() in args.begin_ids and gold[j].item() != args.l2i["[PAD]"]:
                pred_entities_num += 1

        pred_entities_pos, gold_entities_pos = set(), set()

        for j in range(gold.size()[0]):
            if gold[j].item() in args.begin_ids:
                start = j
                for k in range(j+1, gold.size()[0]):
                    if gold[k].item() == args.l2i["[PAD]"] or gold[k].item() == args.l2i["O"] or gold[k].item() in args.begin_ids:
                        end = k - 1
                        break
                else:
                    end = gold.size()[0] - 1
                gold_entities_pos.add((start, end))

        for j in range(pred.size()[0]):
            if pred[j].item() in args.begin_ids and gold[j].item() != args.l2i["[PAD]"]:
                start = j
                for k in range(j+1, pred.size()[0]):
                    if pred[k].item() == args.l2i["[PAD]"] or pred[k].item() == args.l2i["O"] or pred[k].item() in args.begin_ids:
                        end = k - 1
                        break
                else:
                    end = pred.size()[0] - 1
                pred_entities_pos.add((start, end))

        for entity in pred_entities_pos:
            if entity not in gold_entities_pos:
                continue
            for j in range(entity[0], entity[1]+1):
                if gold[j].item() != pred[j].item():
                    break
            else:
                correct += 1

    print("Report precision, recall, and f1:")
    p = correct/pred_entities_num
    r = correct/gold_entities_num
    f1 = 2*p*r/(p+r)
    print("{:.3f}, {:.3f}, {:.3f}".format(p,r,f1))

    return f1


def main():
    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)

    # Path options.
    parser.add_argument("--pretrained_model_path", default=None, type=str,
                        help="Path of the pretrained model.")
    parser.add_argument("--output_model_path", default="./models/ner_model.bin", type=str,
                        help="Path of the output model.")
    parser.add_argument("--vocab_path", default=None, type=str,
                        help="Path of the vocabulary file.")
    parser.add_argument("--spm_model_path", default=None, type=str,
                        help="Path of the sentence piece model.")
    parser.add_argument("--train_path", type=str, required=True,
                        help="Path of the trainset.")
    parser.add_argument("--dev_path", type=str, required=True,
                        help="Path of the devset.")
    parser.add_argument("--test_path", type=str,
                        help="Path of the testset.")
    parser.add_argument("--config_path", default="./models/bert_base_config.json", type=str,
                        help="Path of the config file.")
    parser.add_argument("--label2id_path", type=str, required=True,
                        help="Path of the label2id file.")

    # Model options.
    parser.add_argument("--batch_size", type=int, default=32,
                        help="Batch_size.")
    parser.add_argument("--seq_length", default=128, type=int,
                        help="Sequence length.")
    parser.add_argument("--embedding", choices=["bert", "word"], default="bert",
                        help="Emebdding type.")
    parser.add_argument("--encoder", choices=["bert", "lstm", "gru", \
                                              "cnn", "gatedcnn", "attn", "synt", \
                                              "rcnn", "crnn", "gpt", "bilstm"], \
                                              default="bert", help="Encoder type.")
    parser.add_argument("--bidirectional", action="store_true", help="Specific to recurrent model.")
    parser.add_argument("--factorized_embedding_parameterization", action="store_true", help="Factorized embedding parameterization.")
    parser.add_argument("--parameter_sharing", action="store_true", help="Parameter sharing.")
    
    # Optimizer options.
    parser.add_argument("--learning_rate", type=float, default=2e-5,
                        help="Learning rate.")
    parser.add_argument("--warmup", type=float, default=0.1,
                        help="Warm up value.")
    parser.add_argument("--fp16", action='store_true',
                        help="Whether to use 16-bit (mixed) precision (through NVIDIA apex) instead of 32-bit")
    parser.add_argument("--fp16_opt_level", choices=["O0", "O1", "O2", "O3" ], default='O1',
                        help="For fp16: Apex AMP optimization level selected in ['O0', 'O1', 'O2', and 'O3']."
                             "See details at https://nvidia.github.io/apex/amp.html")

    # Training options.
    parser.add_argument("--dropout", type=float, default=0.1,
                        help="Dropout.")
    parser.add_argument("--epochs_num", type=int, default=3,
                        help="Number of epochs.")
    parser.add_argument("--report_steps", type=int, default=100,
                        help="Specific steps to print prompt.")
    parser.add_argument("--seed", type=int, default=7,
                        help="Random seed.")
    
    args = parser.parse_args()

    # Load the hyperparameters of the config file.
    args = load_hyperparam(args)

    set_seed(args.seed)

    args.begin_ids = []

    with open(args.label2id_path, mode="r", encoding="utf-8") as f:
        l2i = json.load(f)
        print("Labels: ", l2i)
        l2i["[PAD]"] = len(l2i)
        for label in l2i:
            if label.startswith("B"):
                args.begin_ids.append(l2i[label])

    args.l2i = l2i

    args.labels_num = len(l2i)

    args.tokenizer = SpaceTokenizer(args)

    # Build sequence labeling model.
    model = NerTagger(args)

    # Load or initialize parameters.
    load_or_initialize_parameters(args, model)
    
    args.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = model.to(args.device)

    # Training phase.
    instances = read_dataset(args, args.train_path)

    src = torch.LongTensor([ins[0] for ins in instances])
    tgt = torch.LongTensor([ins[1] for ins in instances])
    seg = torch.LongTensor([ins[2] for ins in instances])

    instances_num = src.size(0)
    batch_size = args.batch_size
    args.train_steps = int(instances_num * args.epochs_num / batch_size) + 1

    print("Batch size: ", batch_size)
    print("The number of training instances:", instances_num)

    optimizer, scheduler = build_optimizer(args, model)

    if args.fp16:
        try:
            from apex import amp
        except ImportError:
            raise ImportError("Please install apex from https://www.github.com/nvidia/apex to use fp16 training.")
        model, optimizer = amp.initialize(model, optimizer,opt_level = args.fp16_opt_level)
    
    if torch.cuda.device_count() > 1:
        print("{} GPUs are available. Let's use them.".format(torch.cuda.device_count()))
        model = torch.nn.DataParallel(model)
    args.model = model

    total_loss, f1, best_f1 = 0., 0., 0.

    print("Start training.")

    for epoch in range(1, args.epochs_num+1):
        model.train()
        for i, (src_batch, tgt_batch, seg_batch) in enumerate(batch_loader(batch_size, src, tgt, seg)):
            loss = train(args, model, optimizer, scheduler, src_batch, tgt_batch, seg_batch)
            total_loss += loss.item()
            if (i + 1) % args.report_steps == 0:
                print("Epoch id: {}, Training steps: {}, Avg loss: {:.3f}".format(epoch, i+1, total_loss / args.report_steps))
                total_loss = 0.

        f1 = evaluate(args, read_dataset(args, args.dev_path))
        if f1 > best_f1:
            best_f1 = f1
            save_model(model, args.output_model_path)
        else:
            continue

    # Evaluation phase.
    if args.test_path is not None:
        print("Test set evaluation.")
        if torch.cuda.device_count() > 1:
            model.module.load_state_dict(torch.load(args.output_model_path))
        else:
            model.load_state_dict(torch.load(args.output_model_path))
        evaluate(args, read_dataset(args, args.test_path))


if __name__ == "__main__":
    main()
